import numpy as np
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from sentence_transformers import SentenceTransformer
import os

def get_model_details(model_path: str):
    """
    通过 SentenceTransformer 库获取本地模型的详细信息。
    """
    details = {}
    try:
        # 使用 SentenceTransformer 直接加载模型，以获取底层配置
        sbert_model = SentenceTransformer(model_path)

        # 1. 获取模型配置和参数
        # 模型的配置通常存储在 model[0] (Transformer Model) 的 config 中
        transformer_model = sbert_model._last_module().auto_model

        details['parameters'] = transformer_model.num_parameters()
        details['hidden_size'] = transformer_model.config.hidden_size

        # 2. 获取输出向量维度
        # 直接通过 encode 一个测试文本来获取输出维度
        test_vector = sbert_model.encode("test", convert_to_numpy=True)
        details['embedding_dimension'] = test_vector.shape[0]

        # 3. 获取模型大小（文件大小）
        # 查找 pytorch_model.bin 文件或 model.safetensors 文件的大小
        model_file_path = os.path.join(model_path, "pytorch_model.bin")
        if not os.path.exists(model_file_path):
            # 尝试查找 safetensors 文件
            model_file_path = os.path.join(model_path, "model.safetensors")

        if os.path.exists(model_file_path):
            size_bytes = os.path.getsize(model_file_path)
            # 转换为 GB
            details['size_gb'] = size_bytes / (1024 * 1024 * 1024)
        else:
            details['size_gb'] = "N/A (model file not found)"

    except Exception as e:
        details['error'] = f"无法获取详细信息: {e}"
        details['embedding_dimension'] = "未知"

    return details


# 移动到函数外部
def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))


def test_embedding_consistency():
    # 注意：请确保这些路径下的模型文件是完整的！
    MODEL_PATH = "D:/models/sentence-transformersall-MiniLM-L6-v2"

    print("=========================================")
    print(f"当前测试模型路径: {MODEL_PATH}")

    # 获取并打印模型详细信息
    model_details = get_model_details(MODEL_PATH)
    if 'error' in model_details:
        print(f"⚠️ 警告: {model_details['error']}")
    else:
        print(f"⚙️ 模型维度: {model_details.get('embedding_dimension')} (向量维度)")
        # 参数量通常以 Millions 计算
        params_m = model_details.get('parameters', 0) / 1_000_000
        print(f"⚙️ 模型参数量: {params_m:.2f} M ({model_details.get('hidden_size')} 隐藏层维度)")
        print(f"⚙️ 模型文件大小: {model_details.get('size_gb'):.2f} GB")
    print("=========================================")

    # 初始化 LlamaIndex 嵌入模型
    try:
        embed_model = HuggingFaceEmbedding(model_name=MODEL_PATH)
    except Exception as e:
        print(f"❌ 初始化 HuggingFaceEmbedding 失败，请检查模型文件是否完整: {e}")
        return

    text1 = "人工智能技术"
    text2 = "人工智能技术"  # 完全相同
    text3 = "人工 智能 技术"  # 稍有不同（空格可能会影响某些分词器）
    text4 = "机器学习算法"  # 完全不同

    print("=== 向量一致性测试 ===")

    texts = [text1, text2, text3, text4]
    embeddings = []

    for text in texts:
        # 使用 LlamaIndex 接口获取向量
        embedding = embed_model.get_text_embedding(text)
        embeddings.append(embedding)
        print(f"文本: '{text}'")
        print(f"向量前3维: {embedding[:3]}")

    # --- 相似度分析 ---
    print(f"\n相似度分析 (余弦相似度):")

    # 相同文本：理想值 1.0
    sim_same = cosine_similarity(np.array(embeddings[0]), np.array(embeddings[1]))
    print(f"完全相同文本 (1 vs 2): {sim_same:.6f}")

    # 稍有不同：理想值接近 1.0 (例如 0.85 ~ 0.95)
    sim_minor_diff = cosine_similarity(np.array(embeddings[0]), np.array(embeddings[2]))
    print(f"稍不同文本 (1 vs 3): {sim_minor_diff:.6f}")

    # 完全不同：理想值接近 0.0 或较低（例如 0.4 ~ 0.7）
    sim_diff = cosine_similarity(np.array(embeddings[0]), np.array(embeddings[3]))
    print(f"完全不同文本 (1 vs 4): {sim_diff:.6f}")


test_embedding_consistency()